import requests
import time
import json
import uuid

# 配置信息
BACKEND_URL = "http://localhost:8001"

# 生成UUID格式的session_id
session_id = str(uuid.uuid4())
stream_session_id = str(uuid.uuid4())

# 测试函数
def test_backend_service():
    """测试后端服务是否正常运行"""
    try:
        response = requests.get(f"{BACKEND_URL}/status/system", timeout=5)
        if response.status_code == 200:
            print(f"后端服务正常，状态码: {response.status_code}")
            return True
        else:
            print(f"服务状态检查失败，状态码: {response.status_code}")
            print(f"响应内容: {response.text}")
            return False
    except Exception as e:
        print(f"后端服务连接失败: {str(e)}")
        return False

def test_chat_api():
    """测试聊天接口"""
    try:
        url = f"{BACKEND_URL}/chat"
        payload = {
            "session_id": session_id,
            "message": "你好，测试聊天接口",
            "use_rag": False,
            "use_tools": False,
            "temperature": 0.7,
            "max_tokens": 100
        }
        headers = {"Content-Type": "application/json"}
        
        start_time = time.time()
        response = requests.post(url, json=payload, headers=headers, timeout=30)
        end_time = time.time()
        
        if response.status_code == 200:
            data = response.json()
            print(f"聊天接口测试成功，响应时间: {end_time - start_time:.2f}秒")
            print(f"响应内容: {data['response']}")
            return True
        else:
            print(f"聊天请求失败，状态码: {response.status_code}")
            print(f"响应内容: {response.text}")
            return False
    except Exception as e:
        print(f"聊天接口请求异常: {str(e)}")
        return False

def test_streaming_chat_api():
    """测试流式聊天接口"""
    try:
        url = f"{BACKEND_URL}/chat/stream"
        payload = {
            "session_id": stream_session_id,
            "message": "你好，测试流式聊天接口",
            "use_rag": False,
            "use_tools": False,
            "temperature": 0.7,
            "max_tokens": 100
        }
        headers = {"Content-Type": "application/json"}
        
        start_time = time.time()
        response = requests.post(url, json=payload, headers=headers, stream=True, timeout=60)
        
        if response.status_code == 200:
            full_response = ""
            for line in response.iter_lines(decode_unicode=True):
                if line:
                    # 处理SSE格式
                    if line.startswith('data: '):
                        try:
                            data = json.loads(line[5:])
                            if 'chunk' in data:
                                full_response += data['chunk']
                            elif 'end' in data and data['end']:
                                break
                        except json.JSONDecodeError:
                            pass
            
            end_time = time.time()
            print(f"流式聊天接口测试成功")
            print(f"完整流式响应: {full_response}")
            print(f"完整流式响应长度: {len(full_response)} 字符")
            print(f"流式响应时间: {end_time - start_time:.2f}秒")
            return True
        else:
            print(f"流式聊天请求失败，状态码: {response.status_code}")
            print(f"响应内容: {response.text}")
            return False
    except Exception as e:
        print(f"流式聊天接口请求异常: {str(e)}")
        return False

def test_simple_get_streaming():
    """测试简单的GET流式聊天接口"""
    try:
        url = f"{BACKEND_URL}/chat/stream?q=你好，这是GET流式测试"
        
        start_time = time.time()
        response = requests.get(url, stream=True, timeout=60)
        
        if response.status_code == 200:
            full_response = ""
            for line in response.iter_lines(decode_unicode=True):
                if line:
                    if line.startswith('data: '):
                        # 简单处理，获取实际内容
                        content_line = line[5:]
                        if content_line != '[DONE]':
                            full_response += content_line
            
            end_time = time.time()
            print(f"GET流式聊天接口测试成功")
            print(f"完整GET流式响应: {full_response}")
            print(f"完整GET流式响应长度: {len(full_response)} 字符")
            print(f"GET流式响应时间: {end_time - start_time:.2f}秒")
            return True
        else:
            print(f"GET流式聊天请求失败，状态码: {response.status_code}")
            print(f"响应内容: {response.text}")
            return False
    except Exception as e:
        print(f"GET流式聊天接口请求异常: {str(e)}")
        return False

# 主函数
def main():
    print("=== AI助手系统端到端测试 ===")
    print(f"测试后端服务: {BACKEND_URL}")
    print(f"使用的session_id: {session_id}")
    print()
    
    # 1. 检查后端服务状态
    print("1. 检查后端服务状态...")
    service_ok = test_backend_service()
    print()
    
    # 2. 测试聊天接口
    print("2. 测试聊天接口...")
    chat_ok = test_chat_api()
    print()
    
    # 3. 测试流式聊天接口
    print("3. 测试流式聊天接口...")
    streaming_ok = test_streaming_chat_api()
    print()
    
    # 4. 测试简单的GET流式聊天接口
    print("4. 测试GET流式聊天接口...")
    get_streaming_ok = test_simple_get_streaming()
    print()
    
    # 总结测试结果
    all_tests_passed = service_ok and chat_ok and streaming_ok and get_streaming_ok
    print("=== 端到端测试完成 ===")
    print(f"所有测试通过: {all_tests_passed}")

if __name__ == "__main__":
    main()